#include <ntdll.h>
#include <pe.h>
#include <msg.h>
#include <utils.h>
#include <tlzma.h>
#include <main.h>

#include <ace.h>

HMODULE ace_load_shell_module() {
    wchar_t path[MAX_PATH];
    GetModuleFileNameW(NULL, path, sizeof(path));
    
    wchar_t *extension_dot = wcsrchr(path, L'.');
    wcscpy(extension_dot, L"Base.dll");

    return utils_load_module_patched(path);
}

HMODULE ace_load_runtime_module(const wchar_t *path, const char *symbol, void *target) {
    HMODULE module = utils_load_module_patched(path);
    
    void *symbol_addr = GetProcAddress(module, symbol);
    utils_hook_address(symbol_addr, target);

    return module;
}

void *ace_unpack_shell_data(const char *section_name, HMODULE base_module, size_t *out_size) {
    IMAGE_SECTION_HEADER *section = pe_find_section(base_module, section_name);
    if (!section) {
        msg_err_a("Failed to find the shell data section in the base module. " ISSUE_SUFFIX);
    }

    void *data = ((uint8_t*)base_module) + section->VirtualAddress;
    return tlzma_decode(NULL, data, out_size, NULL);
}

void *ace_unpack_exe_data(size_t *out_size) {
    char *exe = (char*)GetModuleHandleA(NULL);

    IMAGE_DOS_HEADER *dos_header = (IMAGE_DOS_HEADER*)exe;
    IMAGE_NT_HEADERS64 *nt_headers = (IMAGE_NT_HEADERS64*)(exe + dos_header->e_lfanew);

    size_t section_count = nt_headers->FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER *section_headers = (IMAGE_SECTION_HEADER*)(nt_headers + 1);

    IMAGE_SECTION_HEADER *last_section = &section_headers[section_count - 1];
    size_t offset = 0;
    while (offset < last_section->SizeOfRawData && !tlzma_test_header(exe + last_section->VirtualAddress + offset)) {
        offset++;
    }

    if (offset == last_section->SizeOfRawData) {
        msg_err_a("Failed to find exe data blob. " ISSUE_SUFFIX);
    }

    void *data = exe + last_section->VirtualAddress + offset; 
    return tlzma_decode(NULL, data, out_size, NULL);
}

static void ace_unpack_section(void *data, size_t size) {
    void *packed = malloc(size);
    memcpy(packed, data, size);

    uint8_t *unpacked_ptr = (uint8_t*)data;
    uint8_t *packed_ptr = (uint8_t*)packed;

    while (tlzma_test_header(packed_ptr)) {
        size_t unpacked_size, packed_size;
        tlzma_decode(unpacked_ptr, packed_ptr, &unpacked_size, &packed_size);

        unpacked_ptr += unpacked_size;
        packed_ptr += packed_size;
    }

    free(packed);
}

void ace_unpack_sections() {
    uint8_t *module = (uint8_t*)GetModuleHandleA(NULL);

    IMAGE_DOS_HEADER* dos_header = (IMAGE_DOS_HEADER*)module;
    IMAGE_NT_HEADERS64* nt_headers = (IMAGE_NT_HEADERS64*)(((uint8_t*)module) + dos_header->e_lfanew);

    size_t sectionCount = nt_headers->FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER* section_headers = (IMAGE_SECTION_HEADER*)(nt_headers + 1);

    for (size_t i = 0; i < sectionCount; i++) {
        uint8_t *data = ((uint8_t*)module) + section_headers[i].VirtualAddress;

        if (tlzma_test_header(data)) {
            MEMORY_BASIC_INFORMATION memory_info;
            VirtualQuery(data, &memory_info, sizeof(memory_info));
            
            DWORD oldProtect;
            VirtualProtect(data, memory_info.RegionSize, PAGE_READWRITE, &oldProtect);

            ace_unpack_section(data, memory_info.RegionSize);

            VirtualProtect(data, memory_info.RegionSize, oldProtect, &oldProtect);
        }
    }
}
